package com.zlzlib.limiter;

import lombok.RequiredArgsConstructor;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;

/**
 * @DateTime: 2023/5/11 11:10
 * @Author zlz
 * @Version 1.0
 */
@Component
@RequiredArgsConstructor
public class LimiterRedis implements LimiterManager {

    private final StringRedisTemplate stringRedisTemplate;

    private DefaultRedisScript<Long> redisScript;

    @PostConstruct
    public void init() {
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/rateLimiter.lua")));
    }

    @Override
    public boolean tryAccess(LimitReq limitReq, String key) {
        if (key == null) {
            return true;
        }
        List<String> keys = new ArrayList<>();
        keys.add(key);
        Long count = stringRedisTemplate.execute(redisScript, keys,
                String.valueOf(limitReq.limitCount()),
                String.valueOf(limitReq.timeunit().toSeconds(limitReq.timeout())));
        return count == null || count != 0;
    }
}
